from mlxtend import frequent_patterns
import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules
import matplotlib.pyplot as plt
import networkx as nx

# 读取数据
df = pd.read_excel('./associationRules/餐厅数据.xlsx')

# 假设 df 是交易数据列表，需要将其转换为适合 TransactionEncoder 的格式
# 这里需要根据实际数据格式进行调整
transactions = df.values.tolist()

# 打印交易数据，用于调试
print("交易数据：", transactions)

# 将所有数据转换为字符串类型
transactions = [[str(item) for item in transaction] for transaction in transactions]

te = TransactionEncoder()
te_ary = te.fit(transactions).transform(transactions)
df_encoded = pd.DataFrame(te_ary, columns=te.columns_)

# 使用 apriori 进行关联规则挖掘，降低支持度阈值
frequent_itemsets = apriori(df_encoded, min_support=0.01, use_colnames=True)
frequent_itemsets.sort_values(by='support', ascending=False, inplace=True)

# 打印频繁项集，用于调试
print("频繁项集：", frequent_itemsets)

# 选择 2 项集
print(frequent_itemsets[frequent_itemsets.itemsets.apply(lambda x: len(x) == 2)])

# 生成关联规则
rules = association_rules(frequent_itemsets, metric='confidence', min_threshold=0.1)

# 筛选有效规则
effective = rules[
    (rules['lift'] > 1) & (rules['confidence'] > 0.1)
].sort_values(by=['lift', 'confidence'], ascending=False)

# 查看筛选后的规则
print(effective[['antecedents', 'consequents', 'support', 'confidence', 'lift']])

# 设置字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

#生成相应的模型
frequent_itemsets.to_pickle("frequent_itemsets.pkl")
effective.to_pickle("rule.pkl")

# 关联关系可视化
#G = nx.DiGraph()
#for _, row in effective.iterrows():
 #   G.add_edge(','.join(list(row['antecedents'])),
  #             ','.join(list(row['consequents'])),
   #            weight=row['lift'])
#
#plt.figure(figsize=(12, 8))
#pos = nx.spring_layout(G)
#nx.draw(G, pos,
 #       with_labels=True,
  #      edge_color=[G[u][v]['weight'] for u, v in G.edges()],
   #     width=2.0,
    #    edge_cmap=plt.cm.Blues)
#plt.show()